"""
增强版NSGA-II算法实现模块

NSGA-II (Non-dominated Sorting Genetic Algorithm II) 是一种多目标优化算法，
用于解决多个冲突目标的优化问题。

增强版NSGA-II算法包含以下改进：
1. 自适应交叉和变异率
2. 精英保留策略
3. 局部搜索机制
4. 并行评估
5. 早期停止机制
"""

import random
import logging
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed


class EnhancedNSGA2Optimizer:
    """增强版NSGA-II多目标优化算法实现"""

    def __init__(
        self,
        population_size=30,
        crossover_rate=0.5,
        mutation_rate=0.15,
        phenotype_limit=300,
        elimination_size=5,
        max_no_improvement=50,
        adaptive_rates=True,
        elitism=True,
        local_search=True,
        parallel_evaluation=True,
        early_stopping=True,
    ):
        """初始化增强版NSGA-II优化器

        Args:
            population_size: 种群大小
            crossover_rate: 交叉率
            mutation_rate: 变异率
            phenotype_limit: 最大迭代代数
            elimination_size: 每代淘汰数量
            max_no_improvement: 最大无改进代数，用于早期停止
            adaptive_rates: 是否使用自适应交叉和变异率
            elitism: 是否使用精英保留策略
            local_search: 是否使用局部搜索机制
            parallel_evaluation: 是否使用并行评估
            early_stopping: 是否使用早期停止机制
        """
        self.population_size = population_size
        self.initial_crossover_rate = crossover_rate
        self.initial_mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
        self.mutation_rate = mutation_rate
        self.phenotype_limit = phenotype_limit
        self.elimination_size = elimination_size
        self.max_no_improvement = max_no_improvement
        self.adaptive_rates = adaptive_rates
        self.elitism = elitism
        self.local_search = local_search
        self.parallel_evaluation = parallel_evaluation
        self.early_stopping = early_stopping
        self.logger = logging.getLogger(__name__)

    def optimize(
        self,
        initial_solution,
        objective_functions,
        create_variant_func,
        crossover_func,
        mutate_func,
        local_search_func=None,
        context=None,
    ):
        """使用增强版NSGA-II算法优化解决方案

        Args:
            initial_solution: 初始解决方案
            objective_functions: 目标函数列表，每个函数接收一个解决方案并返回一个适应度值
            create_variant_func: 创建变异解决方案的函数，接收一个解决方案和上下文
            crossover_func: 交叉函数，接收两个父代解决方案并返回两个子代
            mutate_func: 变异函数，接收一个解决方案和上下文并返回变异后的解决方案
            local_search_func: 局部搜索函数，接收一个解决方案和上下文并返回改进的解决方案
            context: 上下文信息，传递给create_variant_func和mutate_func

        Returns:
            优化后的解决方案
        """
        self.logger.info("开始使用增强版NSGA-II算法优化解决方案...")
        start_time = datetime.now()

        # 初始化种群
        population = [initial_solution]
        for _ in range(self.population_size - 1):
            variant = create_variant_func(initial_solution, context)
            population.append(variant)

        # 评估初始种群
        fitness_values = self.evaluate_population(population, objective_functions)

        # 计算初始种群的非支配排序和拥挤度
        fronts = self.non_dominated_sort(fitness_values)
        crowding_distances = self.calculate_crowding_distance(fronts, fitness_values)

        # 找到初始最佳解
        best_solution = initial_solution
        best_fitness = self.evaluate_solution(best_solution, objective_functions)
        best_generation = 0
        self.logger.info(f"初始解适应度: {best_fitness}")

        # 记录无改进代数
        no_improvement_count = 0

        # NSGA-II主循环
        for generation in range(self.phenotype_limit):
            # 更新交叉率和变异率
            if self.adaptive_rates:
                self.update_rates(generation)

            # 选择父代
            parents = self.tournament_selection(population, fronts, crowding_distances)

            # 交叉和变异
            offspring = []
            for i in range(0, len(parents), 2):
                if i + 1 < len(parents):
                    if random.random() < self.crossover_rate:
                        child1, child2 = crossover_func(parents[i], parents[i + 1])
                        offspring.extend([child1, child2])
                    else:
                        offspring.extend([parents[i], parents[i + 1]])

            # 变异
            for i in range(len(offspring)):
                if random.random() < self.mutation_rate:
                    offspring[i] = mutate_func(offspring[i], context)

            # 局部搜索
            if self.local_search and local_search_func:
                for i in range(len(offspring)):
                    if random.random() < 0.1:  # 10%的概率进行局部搜索
                        offspring[i] = local_search_func(offspring[i], context)

            # 精英保留
            if self.elitism:
                # 保留上一代的最优解
                elite = [population[fronts[0][0]]]
                offspring = elite + offspring[: len(offspring) - 1]

            # 合并父代和子代
            combined = population + offspring

            # 评估合并种群
            combined_fitness = self.evaluate_population(combined, objective_functions)

            # 非支配排序
            combined_fronts = self.non_dominated_sort(combined_fitness)
            combined_crowding = self.calculate_crowding_distance(
                combined_fronts, combined_fitness
            )

            # 选择下一代
            population = self.select_next_generation(
                combined, combined_fronts, combined_crowding
            )

            # 更新最佳解
            fitness_values = self.evaluate_population(population, objective_functions)
            fronts = self.non_dominated_sort(fitness_values)

            # 找到当前最佳解
            current_best = population[fronts[0][0]]
            current_best_fitness = fitness_values[fronts[0][0]]

            # 更新全局最佳解
            if self.is_better(current_best_fitness, best_fitness):
                best_solution = current_best
                best_fitness = current_best_fitness
                best_generation = generation
                no_improvement_count = 0
                self.logger.info(
                    f"第 {generation} 代发现更好的解，适应度: {best_fitness}"
                )
            else:
                no_improvement_count += 1

            # 早期停止
            if self.early_stopping and no_improvement_count >= self.max_no_improvement:
                self.logger.info(
                    f"连续 {self.max_no_improvement} 代没有改进，提前结束优化"
                )
                break

            # 每10代输出一次日志
            if generation % 10 == 0:
                self.logger.info(
                    f"第 {generation} 代完成，当前最佳适应度: {best_fitness}, 无改进代数: {no_improvement_count}"
                )

        end_time = datetime.now()
        optimization_time = (end_time - start_time).total_seconds()
        self.logger.info(
            f"增强版NSGA-II算法优化完成，耗时 {optimization_time:.2f} 秒，最终适应度: {best_fitness}，最佳解出现在第 {best_generation} 代"
        )

        return best_solution

    def update_rates(self, generation):
        """更新交叉率和变异率

        随着代数的增加，降低交叉率，增加变异率，以增强局部搜索能力

        Args:
            generation: 当前代数
        """
        progress = min(1.0, generation / self.phenotype_limit)
        self.crossover_rate = self.initial_crossover_rate * (1.0 - 0.5 * progress)
        self.mutation_rate = self.initial_mutation_rate * (1.0 + progress)

    def evaluate_population(self, population, objective_functions):
        """评估种群中所有解决方案的适应度

        Args:
            population: 解决方案列表
            objective_functions: 目标函数列表

        Returns:
            适应度值列表，每个元素是一个解决方案的多个目标函数值
        """
        if self.parallel_evaluation:
            return self.parallel_evaluate_population(population, objective_functions)
        else:
            return self.sequential_evaluate_population(population, objective_functions)

    def sequential_evaluate_population(self, population, objective_functions):
        """顺序评估种群

        Args:
            population: 解决方案列表
            objective_functions: 目标函数列表

        Returns:
            适应度值列表
        """
        fitness_values = []
        for solution in population:
            fitness = self.evaluate_solution(solution, objective_functions)
            fitness_values.append(fitness)
        return fitness_values

    def parallel_evaluate_population(self, population, objective_functions):
        """并行评估种群

        Args:
            population: 解决方案列表
            objective_functions: 目标函数列表

        Returns:
            适应度值列表
        """
        fitness_values = [None] * len(population)

        with ThreadPoolExecutor(max_workers=min(8, len(population))) as executor:
            future_to_index = {
                executor.submit(
                    self.evaluate_solution, solution, objective_functions
                ): i
                for i, solution in enumerate(population)
            }

            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    fitness_values[index] = future.result()
                except Exception as e:
                    self.logger.error(f"评估解决方案时出错: {str(e)}")
                    # 使用一个默认的低适应度值
                    fitness_values[index] = [0.0] * len(objective_functions)

        return fitness_values

    def evaluate_solution(self, solution, objective_functions):
        """评估单个解决方案的适应度

        Args:
            solution: 解决方案
            objective_functions: 目标函数列表

        Returns:
            目标函数值列表
        """
        return [func(solution) for func in objective_functions]

    def is_better(self, fitness1, fitness2):
        """判断fitness1是否优于fitness2

        Args:
            fitness1: 第一个适应度值列表
            fitness2: 第二个适应度值列表

        Returns:
            如果fitness1优于fitness2，则返回True
        """
        better_in_any = False
        for f1, f2 in zip(fitness1, fitness2):
            if f1 < f2:  # 假设所有目标都是最小化问题
                return False
            if f1 > f2:
                better_in_any = True
        return better_in_any

    def non_dominated_sort(self, fitness_values):
        """非支配排序

        Args:
            fitness_values: 适应度值列表

        Returns:
            前沿列表，每个前沿是一个解决方案索引列表
        """
        n = len(fitness_values)
        domination_count = [0] * n  # 被支配的次数
        dominated_solutions = [[] for _ in range(n)]  # 支配的解决方案
        fronts = [[]]  # 前沿

        # 计算支配关系
        for i in range(n):
            for j in range(n):
                if i == j:
                    continue

                # 检查i是否支配j
                if self.dominates(fitness_values[i], fitness_values[j]):
                    dominated_solutions[i].append(j)
                # 检查j是否支配i
                elif self.dominates(fitness_values[j], fitness_values[i]):
                    domination_count[i] += 1

            # 如果i不被任何解支配，则属于第一前沿
            if domination_count[i] == 0:
                fronts[0].append(i)

        # 计算其他前沿
        i = 0
        while fronts[i]:
            next_front = []
            for j in fronts[i]:
                for k in dominated_solutions[j]:
                    domination_count[k] -= 1
                    if domination_count[k] == 0:
                        next_front.append(k)
            i += 1
            if next_front:
                fronts.append(next_front)

        return fronts

    def dominates(self, fitness1, fitness2):
        """判断fitness1是否支配fitness2

        Args:
            fitness1: 第一个适应度值列表
            fitness2: 第二个适应度值列表

        Returns:
            如果fitness1支配fitness2，则返回True
        """
        better_in_any = False
        for f1, f2 in zip(fitness1, fitness2):
            if f1 < f2:  # 假设所有目标都是最小化问题
                return False
            if f1 > f2:
                better_in_any = True
        return better_in_any

    def calculate_crowding_distance(self, fronts, fitness_values):
        """计算拥挤度距离

        Args:
            fronts: 前沿列表
            fitness_values: 适应度值列表

        Returns:
            拥挤度距离列表
        """
        n = len(fitness_values)
        crowding_distance = [0] * n

        for front in fronts:
            if len(front) <= 2:
                for i in front:
                    crowding_distance[i] = float("inf")
                continue

            # 对每个目标函数计算拥挤度
            for m in range(len(fitness_values[0])):
                # 按目标函数值排序
                sorted_front = sorted(front, key=lambda i: fitness_values[i][m])

                # 边界点的拥挤度设为无穷大
                crowding_distance[sorted_front[0]] = float("inf")
                crowding_distance[sorted_front[-1]] = float("inf")

                # 计算中间点的拥挤度
                f_max = fitness_values[sorted_front[-1]][m]
                f_min = fitness_values[sorted_front[0]][m]

                if f_max == f_min:
                    continue

                for i in range(1, len(sorted_front) - 1):
                    crowding_distance[sorted_front[i]] += (
                        fitness_values[sorted_front[i + 1]][m]
                        - fitness_values[sorted_front[i - 1]][m]
                    ) / (f_max - f_min)

        return crowding_distance

    def tournament_selection(
        self, population, fronts, crowding_distances, tournament_size=2
    ):
        """锦标赛选择

        Args:
            population: 种群
            fronts: 前沿列表
            crowding_distances: 拥挤度距离列表
            tournament_size: 锦标赛大小

        Returns:
            选中的父代列表
        """
        n = len(population)
        selected = []

        # 创建排名列表
        rank = [0] * n
        for i, front in enumerate(fronts):
            for j in front:
                rank[j] = i

        while len(selected) < n:
            # 随机选择tournament_size个个体
            candidates = random.sample(range(n), tournament_size)

            # 选择排名最高的个体
            best_candidate = candidates[0]
            for candidate in candidates[1:]:
                # 如果候选者的排名更好，或者排名相同但拥挤度更大，则选择候选者
                if rank[candidate] < rank[best_candidate] or (
                    rank[candidate] == rank[best_candidate]
                    and crowding_distances[candidate]
                    > crowding_distances[best_candidate]
                ):
                    best_candidate = candidate

            selected.append(population[best_candidate])

        return selected

    def select_next_generation(self, combined, fronts, crowding_distances):
        """选择下一代种群

        Args:
            combined: 合并的父代和子代
            fronts: 前沿列表
            crowding_distances: 拥挤度距离列表

        Returns:
            下一代种群
        """
        next_generation = []
        i = 0

        # 按前沿顺序添加解决方案
        while len(next_generation) + len(fronts[i]) <= self.population_size:
            for j in fronts[i]:
                next_generation.append(combined[j])
            i += 1

            # 如果所有前沿都已添加，则返回
            if i >= len(fronts):
                return next_generation

        # 对最后一个要添加的前沿按拥挤度排序
        last_front = sorted(
            fronts[i], key=lambda j: crowding_distances[j], reverse=True
        )

        # 添加拥挤度最大的解决方案，直到种群大小达到要求
        j = 0
        while len(next_generation) < self.population_size:
            next_generation.append(combined[last_front[j]])
            j += 1

        return next_generation
